import numpy as np
import matplotlib.pyplot as plt

# 1. 读取数据集
def load_data(filename):
    with open(filename, 'r') as file:
        data = file.readlines()
    dataset = np.array([list(map(float, d.strip().split(','))) for d in data])
    return dataset

# 2. 实现k-means算法
def kmeans(X, k, max_iters=100):
    m, n = X.shape
    centroids = X[np.random.choice(m, k, replace=False)]
    for _ in range(max_iters):
        C = np.zeros(m)
        for i in range(m):
            distances = np.sqrt(np.sum((X[i] - centroids) ** 2, axis=1))
            C[i] = np.argmin(distances)
        new_centroids = np.array([X[C == j].mean(axis=0) for j in range(k)])
        if np.all(centroids == new_centroids):
            break
        centroids = new_centroids
    return C, centroids

# 3. 计算轮廓系数
def silhouette_score(X, C):
    m, n = X.shape
    silhouette_vals = np.zeros(m)
    for i in range(m):
        cluster = C[i]
        a = np.mean(np.sqrt(np.sum((X[i] - X[C == cluster]) ** 2, axis=1)))
        p = np.min([np.mean(np.sqrt(np.sum((X[i] - X[C == j]) ** 2, axis=1))) for j in range(k) if j != cluster])
        s = (p - a) / max(a, p)
        silhouette_vals[i] = s
    return silhouette_vals

# 4. 画出轮廓图
def plot_silhouette(X, C, silhouette_vals):
    plt.figure(figsize=(10, 7))
    y_lower = 10
    for i in range(k):
        ith_cluster_silhouette_values = silhouette_vals[C == i]
        ith_cluster_silhouette_values.sort()
        size_cluster_i = ith_cluster_silhouette_values.shape[0]
        y_upper = y_lower + size_cluster_i

        color = plt.cm.nipy_spectral(float(i) / k)
        plt.fill_betweenx(np.arange(y_lower, y_upper), 0, ith_cluster_silhouette_values,
                         facecolor=color, edgecolor=color, alpha=0.7)
        plt.text(-0.05, y_lower + 0.5 * size_cluster_i, str(i))

        y_lower = y_upper + 10

    plt.title('The silhouette plot for the various clusters.')
    yticks = np.arange(-1, k + 1, 1)
    yticklabels = map(str, yticks)
    plt.yticks(yticks, yticklabels)
    plt.xlabel('The silhouette coefficient values')
    plt.ylabel('Cluster label')

    plt.show()

# 主程序
if __name__ == '__main__':
    dataset = load_data('./wine/wine.data')
    k = 3  # 假设聚类数为3
    C, centroids = kmeans(dataset, k)
    silhouette_vals = silhouette_score(dataset, C)
    score = np.mean(silhouette_vals)
    print(f'Silhouette Score: {score}')
    plot_silhouette(dataset, C, silhouette_vals)